import numpy as np
import math
from copy import deepcopy
import matplotlib.pyplot as plt

'''
白鲸优化算法
参数
n：白鲸种群大小
tmax：最大迭代次数
lb：变量下限
ub：变量上限
nd：变量维数
fobj：目标函数
'''


# 白鲸优化算法
def bwo(n, tmax, lb, ub, nd, fobj):
    # 记录每轮更新最优解
    best_fval = []
    # 位置矩阵
    x = lb + np.random.random((n, nd)) * (ub - lb)
    # 适应值矩阵
    fobj_list = [fobj(x[i, ...]) for i in range(n)]
    fx = [f() for f in fobj_list]
    # 适应值拷贝矩阵
    temp_fx = deepcopy(fx)
    # 最优解和最优解值
    xposbest = x[np.argmin(fx)]
    fvalbest = min(fx)
    best_fval.append(fvalbest)
    # 循环更新
    for t in range(1, tmax + 1):
        # 位置拷贝矩阵
        temp_x = deepcopy(x)
        # 鲸落概率
        wf = 0.1 - 0.05 * (t / tmax)
        # 计算每只白鲸的平衡因子bf
        b0 = np.random.random(n)
        for b0_index in range(len(b0)):
            while b0[b0_index] == 0:
                b0[b0_index] = np.random.random()
        bf = b0 * (1 - t / (2 * tmax))
        #  更新每一只白鲸位置
        for i in range(n):
            # 探索阶段
            if bf[i] > 0.5:
                r1 = np.random.random()
                r2 = np.random.random()
                while r1 == 0:
                    r1 = np.random.random()
                while r2 == 0:
                    r2 = np.random.random()
                # 选择随机白鲸
                r = np.random.randint(0, n)
                while r == i:
                    r = np.random.randint(0, n)
                pj = np.arange(nd)
                np.random.shuffle(pj)
                if nd <= n / 5:
                    temp_x[i, pj[0]] = x[i, pj[0]] + (x[r, pj[0]] - x[i, pj[0]]) * (1 + r1) * np.sin(2 * np.pi * r2)
                    temp_x[i, pj[1]] = x[i, pj[1]] + (x[r, pj[1]] - x[i, pj[1]]) * (1 + r1) * np.cos(2 * np.pi * r2)
                else:
                    for j in range(int(nd / 2)):
                        temp_x[i, 2 * j] = x[i, 2 * j] + (x[r, pj[1]] - x[i, 2 * j]) * (1 + r1) * np.sin(2 * np.pi * r2)
                        temp_x[i, 2 * j + 1] = x[i, 2 * j + 1] + (x[r, pj[1]] - x[i, 2 * j + 1]) * (1 + r1) * np.cos(
                            2 * np.pi * r2)
            # 开发阶段
            else:
                r3 = np.random.random()
                r4 = np.random.random()
                while r3 == 0:
                    r3 = np.random.random()
                while r4 == 0:
                    r4 = np.random.random()
                c1 = 2 * r4 * (1 - (t / tmax))
                # 随机白鲸
                r = np.random.randint(0, n)
                while r == i:
                    r = np.random.randint(0, n)
                beta = 3 / 2
                delta = ((math.gamma(1 + beta) * np.sin((np.pi * beta) / 2))
                         / (math.gamma((1 + beta) / 2) * beta * (2 ** ((beta - 1) / 2)))) ** (1 / beta)
                u = np.random.randn(nd)
                v = np.random.randn(nd)
                lf = 0.05 * ((u * delta) / (abs(v) ** (1 / beta)))
                temp_x[i, ...] = r3 * xposbest - r4 * x[i, ...] + c1 * lf * (x[r, ...] - x[i, ...])
            # 处理超出边界值
            flaglb = np.asarray(temp_x[i, ...] < lb).astype(np.int8)
            flagub = np.asarray(temp_x[i, ...] > ub).astype(np.int8)
            flag = flaglb + flagub
            for k in range(nd):
                if flag[k] == 0:
                    flag[k] = 1
                else:
                    flag[k] = 0
            temp_x[i, ...] = temp_x[i, ...] * flag + ub * flagub + lb * flaglb
            result = fobj(temp_x[i, ...])
            temp_fx[i] = result()
            if temp_fx[i] < fx[i]:
                x[i, ...] = temp_x[i, ...]
                fx[i] = temp_fx[i]
        # 鲸落
        for i in range(n):
            if bf[i] <= wf:
                r5 = np.random.random()
                r6 = np.random.random()
                r7 = np.random.random()
                while r5 == 0:
                    r5 = np.random.random()
                while r6 == 0:
                    r6 = np.random.random()
                while r7 == 0:
                    r7 = np.random.random()
                c2 = 2 * wf * n
                # 随机白鲸
                r = np.random.randint(0, n)
                xstep = (ub - lb) * (np.exp((-c2 * t) / tmax))
                temp_x[i, ...] = r5 * x[i, ...] - r6 * x[r, ...] + r7 * xstep
                # 处理超出边界值
                flaglb = np.asarray(temp_x[i, ...] < lb).astype(np.int8)
                flagub = np.asarray(temp_x[i, ...] > ub).astype(np.int8)
                flag = flaglb + flagub
                for m in range(nd):
                    if flag[m] == 0:
                        flag[m] = 1
                    else:
                        flag[m] = 0
                temp_x[i, ...] = temp_x[i, ...] * flag + ub * flagub + lb * flaglb
                result = fobj(temp_x[i, ...])
                temp_fx[i] = result()
                if temp_fx[i] < fx[i]:
                    x[i, ...] = temp_x[i, ...]
                    fx[i] = temp_fx[i]
        # 更改最优值
        temp_xposbest = x[np.argmin(fx)]
        temp_fvalbest = min(fx)
        if temp_fvalbest < fvalbest:
            fvalbest = temp_fvalbest
            xposbest = temp_xposbest
        best_fval.append(fvalbest)
    return xposbest, fvalbest, best_fval


# mu函数
def mu_fun(x, a, k, m):
    if x > a:
        return k * ((x - a) ** m)
    elif -a <= x <= a:
        return 0
    elif x < a:
        return k * ((-x - a) ** m)


# 求解适应值
def get_function_result(fname, *args):
    result = -1
    if fname == "f1":
        def f1():
            return sum(args[0] ** 2)

        result = f1
    if fname == "f2":
        def f2():
            sum1 = 0
            sum2 = 1
            for i in range(len(args[0])):
                sum1 += abs(args[0][i])
                sum2 *= abs(args[0][i])
            return sum1 + sum2

        result = f2
    elif fname == "f3":
        def f3():
            sum1 = 0
            for i in range(len(args[0])):
                sum1 += abs(args[0][i]) ** (i + 1)
            return sum1

        result = f3
    elif fname == "f4":
        def f4():
            sum1 = 0
            sum2 = 0
            for i in range(len(args[0])):
                sum1 += args[0][i]
            for j in range(len(args[0])):
                sum2 += sum1 ** 2
            return sum2

        result = f4
    elif fname == "f5":
        def f5():
            return np.max(abs(args[0]))

        result = f5
    elif fname == "f6":
        def f6():
            sum1 = 0
            for i in range(len(args[0]) - 1):
                sum1 += (100 * ((args[0][i + 1] - (args[0][i] ** 2)) ** 2) + ((args[0][i] - 1) ** 2))
            return sum1

        result = f6
    elif fname == "f7":
        def f7():
            sum1 = 0
            for i in range(len(args[0])):
                sum1 += ((args[0][i] + 0.5) ** 2)
            return sum1

        result = f7
    elif fname == "f8":
        def f8():
            sum1 = 0
            for i in range(len(args[0])):
                sum1 += (i + 1) * (args[0][i] ** 4)
            return sum1 + np.random.random()

        result = f8
    elif fname == "f9":
        def f9():
            sum1 = 0
            sum2 = 0
            for i in range(len(args[0])):
                sum1 += args[0][i] ** 2
                sum2 += 0.5 * (i + 1) * args[0][i]
            return sum1 + sum2 ** 2 + sum2 ** 4

        result = f9
    elif fname == "f10":
        def f10():
            sum1 = 0
            for i in range(len(args[0])):
                sum1 += args[0][i] * np.sin(np.sqrt(abs(args[0][i])))
            return -sum1

        result = f10
    elif fname == "f11":
        def f11():
            sum1 = 0
            sum2 = 0
            for i in range(len(args[0])):
                sum1 += np.sin(args[0][i]) ** 2
                sum2 += args[0][i] ** 2
            return 1 + sum1 - np.exp(-sum2)

        result = f11
    elif fname == "f12":
        def f12():
            sum1 = 0
            for i in range(len(args[0])):
                sum1 += (args[0][i] ** 4) - 16 * (args[0][i] ** 2) + 5 * args[0][i]
            return 0.5 * sum1

        result = f12
    elif fname == "f13":
        def f13():
            sum1 = 0
            for i in range(len(args[0])):
                sum1 += (args[0][i] ** 2) - 10 * np.cos(2 * np.pi * args[0][i]) + 10
            return sum1

        result = f13
    elif fname == "f14":
        def f14():
            sum1 = 0
            sum2 = 0
            for i in range(len(args[0])):
                sum1 += args[0][i] ** 2
                sum2 += np.cos(2 * np.pi * args[0][i])
            return -20 * np.exp(-0.2 * np.sqrt(sum1 / len(args[0]))) - np.exp(sum2 / len(args[0])) + 20 + np.e

        result = f14
    elif fname == "f15":
        def f15():
            sum1 = 0
            sum2 = 1
            for i in range(len(args[0])):
                sum1 += args[0][i] ** 2
                sum2 *= np.cos(args[0][i] / np.sqrt(i + 1))
            return sum1 / 4000 - sum2 + 1

        result = f15
    elif fname == "f16":
        def f16():
            sum1 = 0
            sum2 = 0
            sum3 = 0
            for i in range(len(args[0])):
                sum1 += np.sin(args[0][i]) ** 2
                sum2 += args[0][i] ** 2
                sum3 += np.sin(np.sqrt(abs(args[0][i]))) ** 2
            return (sum1 - np.exp(-sum2)) * np.exp(-sum3)

        result = f16
    elif fname == "f17":
        def f17():
            sum1 = 0
            sum2 = 0
            for i in range(len(args[0]) - 1):
                sum1 += ((0.25 * (args[0][i] + 1)) ** 2) * (
                        1 + 10 * (np.sin(np.pi * (1 + 0.25 * (args[0][i + 1] + 1))) ** 2))
            for j in range(len(args[0])):
                sum2 += mu_fun(args[0][j], 10, 100, 4)
            return np.pi / len(args[0]) * (10 * (np.sin(np.pi * (1 + 0.25 * (args[0][0] + 1))) ** 2) + sum1 + (
                    0.25 * (args[0][-1] + 1)) ** 2) + sum2

        result = f17
    elif fname == "f18":
        def f18():
            sum1 = 0
            sum2 = 0
            for i in range(len(args[0]) - 1):
                sum1 += ((args[0][i] - 1) ** 2) * (1 + np.sin(3 * np.pi * args[0][i + 1]) ** 2)
            for j in range(len(args[0])):
                sum2 += mu_fun(args[0][j], 5, 100, 4)
            return 0.1 * (np.sin(3 * np.pi * args[0][0]) ** 2 + sum1 + ((args[0][-1] - 1) ** 2) * (
                    1 + np.sin(2 * np.pi * args[0][-1]) ** 2)) + sum2

        result = f18
    elif fname == "f19":
        def f19():
            sum2 = 0
            a = [[-32, -16, 0, 16, 32, -32, -16, 0, 16, 32, -32, -16, 0, 16, 32, -32, -16, 0, 16, 32, -32, -16, 0, 16,
                  32],
                 [-32, -32, -32, -32, -32, -16, -16, -16, -16, -16, 0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 32, 32, 32, 32,
                  32]]
            for i in range(25):
                sum1 = 0
                for j in range(len(args[0])):
                    sum1 += (args[0][j] - a[j][i]) ** 6
                sum2 += 1 / ((i + 1) + sum1)
            return 1 / (1 / 500 + sum2)

        result = f19
    elif fname == "f20":
        def f20():
            sum1 = 0
            a = [0.1957, 0.1947, 0.1735, 0.16, 0.0844, 0.0627, 0.0456, 0.0342, 0.0323, 0.0235, 0.0246]
            b = 1 / np.asarray([0.25, 0.5, 1, 2, 4, 6, 8, 10, 12, 14, 16])
            for i in range(11):
                sum1 += abs(a[i] - (args[0][0] * (b[i] ** 2 + b[i] * args[0][1])) / (
                        b[i] ** 2 + b[i] * args[0][2] + args[0][3])) ** 2
            return sum1

        result = f20
    elif fname == "f21":
        def f21():
            sum1 = 4 * (args[0][0] ** 2) - 2.1 * (args[0][0] ** 4) + (args[0][0] ** 6) / 3 + args[0][0] * args[0][
                1] - 4 * (args[0][1] ** 2) + 4 * (args[0][1] ** 4)
            return sum1

        result = f21
    elif fname == "f22":
        def f22():
            sum1 = 0
            a = [[4, 4, 4, 4], [1, 1, 1, 1], [8, 8, 8, 8], [6, 6, 6, 6], [3, 7, 3, 7], [2, 9, 2, 9], [5, 5, 3, 3],
                 [8, 1, 8, 1], [6, 2, 6, 2], [7, 3.6, 7, 3.6]]
            c = [0.1, 0.2, 0.2, 0.4, 0.4, 0.6, 0.3, 0.7, 0.5, 0.5]
            for i in range(5):
                sum1 += 1 / abs(sum((args[0] - np.asarray(a[i])) ** 2) + c[i])
            return -sum1

        result = f22
    elif fname == "f23":
        def f23():
            sum1 = 0
            a = [[4, 4, 4, 4], [1, 1, 1, 1], [8, 8, 8, 8], [6, 6, 6, 6], [3, 7, 3, 7], [2, 9, 2, 9], [5, 5, 3, 3],
                 [8, 1, 8, 1], [6, 2, 6, 2], [7, 3.6, 7, 3.6]]
            c = [0.1, 0.2, 0.2, 0.4, 0.4, 0.6, 0.3, 0.7, 0.5, 0.5]
            for i in range(7):
                sum1 += 1 / abs(sum((args[0] - np.asarray(a[i])) ** 2) + c[i])
            return -sum1

        result = f23
    elif fname == "f24":
        def f24():
            sum1 = 0
            a = [[4, 4, 4, 4], [1, 1, 1, 1], [8, 8, 8, 8], [6, 6, 6, 6], [3, 7, 3, 7], [2, 9, 2, 9], [5, 5, 3, 3],
                 [8, 1, 8, 1], [6, 2, 6, 2], [7, 3.6, 7, 3.6]]
            c = [0.1, 0.2, 0.2, 0.4, 0.4, 0.6, 0.3, 0.7, 0.5, 0.5]
            for i in range(10):
                sum1 += 1 / abs(sum((args[0] - np.asarray(a[i])) ** 2) + c[i])
            return -sum1

        result = f24
    return result


# 获取目标函数相关参数
def get_functions_details(fname):
    if fname == "f1":
        lb = -100
        ub = 100
        nd = 30
        fobj = lambda x: get_function_result("f1", x)
        return lb, ub, nd, fobj
    elif fname == "f2":
        lb = -10
        ub = 10
        nd = 30
        fobj = lambda x: get_function_result("f2", x)
        return lb, ub, nd, fobj
    elif fname == "f3":
        lb = -1
        ub = 1
        nd = 30
        fobj = lambda x: get_function_result("f3", x)
        return lb, ub, nd, fobj
    elif fname == "f4":
        lb = -100
        ub = 100
        nd = 30
        fobj = lambda x: get_function_result("f4", x)
        return lb, ub, nd, fobj
    elif fname == "f5":
        lb = -100
        ub = 100
        nd = 30
        fobj = lambda x: get_function_result("f5", x)
        return lb, ub, nd, fobj
    elif fname == "f6":
        lb = -30
        ub = 30
        nd = 30
        fobj = lambda x: get_function_result("f6", x)
        return lb, ub, nd, fobj
    elif fname == "f7":
        lb = -100
        ub = 100
        nd = 30
        fobj = lambda x: get_function_result("f7", x)
        return lb, ub, nd, fobj
    elif fname == "f8":
        lb = -1.28
        ub = 1.28
        nd = 30
        fobj = lambda x: get_function_result("f8", x)
        return lb, ub, nd, fobj
    elif fname == "f9":
        lb = -5
        ub = 10
        nd = 30
        fobj = lambda x: get_function_result("f8", x)
        return lb, ub, nd, fobj
    elif fname == "f10":
        lb = -500
        ub = 500
        nd = 30
        fobj = lambda x: get_function_result("f10", x)
        return lb, ub, nd, fobj
    elif fname == "f11":
        lb = -10
        ub = 10
        nd = 30
        fobj = lambda x: get_function_result("f11", x)
        return lb, ub, nd, fobj
    elif fname == "f12":
        lb = -5
        ub = 5
        nd = 30
        fobj = lambda x: get_function_result("f12", x)
        return lb, ub, nd, fobj
    elif fname == "f13":
        lb = -5.12
        ub = 5.12
        nd = 30
        fobj = lambda x: get_function_result("f13", x)
        return lb, ub, nd, fobj
    elif fname == "f14":
        lb = -32
        ub = 32
        nd = 30
        fobj = lambda x: get_function_result("f14", x)
        return lb, ub, nd, fobj
    elif fname == "f15":
        lb = -600
        ub = 600
        nd = 30
        fobj = lambda x: get_function_result("f15", x)
        return lb, ub, nd, fobj
    elif fname == "f16":
        lb = -10
        ub = 10
        nd = 30
        fobj = lambda x: get_function_result("f16", x)
        return lb, ub, nd, fobj
    elif fname == "f17":
        lb = -50
        ub = 50
        nd = 30
        fobj = lambda x: get_function_result("f17", x)
        return lb, ub, nd, fobj
    elif fname == "f18":
        lb = -50
        ub = 50
        nd = 30
        fobj = lambda x: get_function_result("f18", x)
        return lb, ub, nd, fobj
    elif fname == "f19":
        lb = -65.536
        ub = 65.536
        nd = 2
        fobj = lambda x: get_function_result("f19", x)
        return lb, ub, nd, fobj
    elif fname == "f20":
        lb = -5
        ub = 5
        nd = 4
        fobj = lambda x: get_function_result("f20", x)
        return lb, ub, nd, fobj
    elif fname == "f21":
        lb = -5
        ub = 5
        nd = 2
        fobj = lambda x: get_function_result("f21", x)
        return lb, ub, nd, fobj
    elif fname == "f22":
        lb = 0
        ub = 10
        nd = 4
        fobj = lambda x: get_function_result("f22", x)
        return lb, ub, nd, fobj
    elif fname == "f23":
        lb = 0
        ub = 10
        nd = 4
        fobj = lambda x: get_function_result("f23", x)
        return lb, ub, nd, fobj
    elif fname == "f24":
        lb = 0
        ub = 10
        nd = 4
        fobj = lambda x: get_function_result("f24", x)
        return lb, ub, nd, fobj


if __name__ == '__main__':
    # 白鲸种群数
    out_n = 50
    # 最大迭代次数
    out_tmax = 1000
    # 目标函数
    out_fobj = "f22"
    # 获取函数参数
    out_lb, out_ub, out_nd, out_fobj = get_functions_details(out_fobj)
    out_xposbest, out_fvalbest, out_best_fval = bwo(out_n, out_tmax, out_lb, out_ub, out_nd, out_fobj)
    # 画图
    plt.plot(range(len(out_best_fval)), out_best_fval)
    plt.xlabel("t")
    plt.ylabel("best value")
    plt.show()
    print("xposbest = ", out_xposbest)
    print("fvalbest = ", f"{out_fvalbest: .5f}")
